{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "71293c1d",
   "metadata": {},
   "source": [
    "# Online Banking Assistant\n",
    "\n",
    "This application is the Bank of Transformers online banking application.\n",
    "\n",
    "Using this application users can:\n",
    "\n",
    "- open an account\n",
    "- add funds to an account\n",
    "- withdraw funds from an account\n",
    "- check the balance of accounts\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "989a6bf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from dotenv import load_dotenv\n",
    "from openai import OpenAI\n",
    "import gradio as gr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "94e7f3d1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OpenAI API Key exists and begins sk-proj-\n"
     ]
    }
   ],
   "source": [
    "# Initialization\n",
    "\n",
    "load_dotenv(override=True)\n",
    "\n",
    "openai_api_key = os.getenv('OPENAI_API_KEY')\n",
    "if openai_api_key:\n",
    "    print(f\"OpenAI API Key exists and begins {openai_api_key[:8]}\")\n",
    "else:\n",
    "    print(\"OpenAI API Key not set\")\n",
    "    \n",
    "MODEL = \"gpt-4.1-mini\"\n",
    "openai = OpenAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "d53e4dd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "WELCOME = \"\"\"\\\n",
    "**Welcome to the Bank of Transformers Online Banking App**\n",
    "\n",
    "You can do the following:\n",
    "- open an account\n",
    "- deposit funds\n",
    "- withdraw funds\n",
    "- check the balance of your account\n",
    "\n",
    "What would you like to do today?\n",
    "\"\"\"\n",
    "\n",
    "SYS_PROMPT = \"\"\"\\\n",
    "You are a helpful banking assistant for an online banking app. Your job:\n",
    "1) Understand the user's intent (open account, deposit, withdraw funds, check balance).\n",
    "2) If you need to take an action, use the provided tools (functions).\n",
    "3) Be concise, ask only for the minimum needed info next.\n",
    "4) Always reflect current balances when relevant.\n",
    "5) Never invent accounts or balances—rely on tool results only.\n",
    "6) After carrying out every action for the user, list all available options, as user might want to deposit/withdraw more funds \n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "36d5f7cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "open_account_function = {\n",
    "    \"name\": \"open_account\",\n",
    "    \"description\": \"Open an account for the user. Ask the user to enter their name and email address.\",\n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"name\": {\n",
    "                \"type\": \"string\",\n",
    "                \"description\": \"The name of the user\",\n",
    "            },\n",
    "            \"email\": {\n",
    "                \"type\": \"string\",\n",
    "                \"description\": \"The email address of the user\",\n",
    "            },\n",
    "        },\n",
    "        \"required\": [\"name\", \"email\"],\n",
    "        \"additionalProperties\": False\n",
    "    }\n",
    "}\n",
    "\n",
    "deposit_funds_function = {\n",
    "    \"name\": \"deposit_funds\",\n",
    "    \"description\": \"Deposit funds into the user account. Ask the user to enter the amount to deposit.\",\n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"amount\": {\n",
    "                \"type\": \"number\",\n",
    "                \"description\": \"The amount to deposit\",\n",
    "            },\n",
    "            \"balance\": {\n",
    "                \"type\": \"number\",\n",
    "                \"description\": \"The account balance before the deposit\",\n",
    "            }\n",
    "        },\n",
    "        \"required\": [\"amount\", \"balance\"],\n",
    "        \"additionalProperties\": False\n",
    "    }\n",
    "}\n",
    "\n",
    "withdraw_funds_function = {\n",
    "    \"name\": \"withdraw_funds\",\n",
    "    \"description\": \"Withdraw funds from the user account. Ask the user to enter the amount to withdraw.\",\n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"amount\": {\n",
    "                \"type\": \"number\",\n",
    "                \"description\": \"The amount to withdraw\",\n",
    "            },\n",
    "            \"balance\": {\n",
    "                \"type\": \"number\",\n",
    "                \"description\": \"The account balance before the withdraw\",\n",
    "            }\n",
    "        },\n",
    "        \"required\": [\"amount\", \"balance\"],\n",
    "        \"additionalProperties\": False\n",
    "    }\n",
    "}\n",
    "\n",
    "check_balance_function = {\n",
    "    \"name\": \"check_balance\",\n",
    "    \"description\": \"Get the balance of the user account and display it to the user.\",\n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"balance\": {\n",
    "                \"type\": \"number\",\n",
    "                \"description\": \"The account balance\",\n",
    "            }\n",
    "        },\n",
    "        \"required\": [\"balance\"],\n",
    "        \"additionalProperties\": False\n",
    "    }    \n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "e35f7074",
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = [\n",
    "    {\"type\": \"function\", \"function\": open_account_function}, \n",
    "    {\"type\": \"function\", \"function\": deposit_funds_function},\n",
    "    {\"type\": \"function\", \"function\": withdraw_funds_function},\n",
    "    {\"type\": \"function\", \"function\": check_balance_function}\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "12cb2f44",
   "metadata": {},
   "outputs": [],
   "source": [
    "def open_account(name, email):\n",
    "    return {\"message\": f\"Account opened for {name} with email {email}\"}\n",
    "\n",
    "def deposit_funds(amount, balance):\n",
    "    new_balance = float(balance) + float(amount)\n",
    "    return {\n",
    "        \"message\": f\"Deposited ${amount:.2f} into your account. Balance: ${new_balance:.2f}\",\n",
    "        \"balance\": new_balance\n",
    "    }\n",
    "\n",
    "def withdraw_funds(amount, balance):\n",
    "    new_balance = float(balance) - float(amount)\n",
    "    return {\n",
    "        \"message\": f\"Withdrew ${amount:.2f} from your account. Balance: ${new_balance:.2f}\",\n",
    "        \"balance\": new_balance\n",
    "    }\n",
    "\n",
    "def check_balance(balance):\n",
    "    bal = float(balance)\n",
    "    return {\"message\": f\"Your account balance is: ${bal:.2f}\", \"balance\": bal}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "c05ada4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We have to write that function handle_tool_call:\n",
    "\n",
    "def handle_tool_call(assistant_msg, balance):\n",
    "    tool_messages = []\n",
    "    new_balance = balance\n",
    "\n",
    "    for tc in getattr(assistant_msg, \"tool_calls\", []) or []:\n",
    "        args = json.loads(tc.function.arguments or \"{}\")\n",
    "        name = tc.function.name\n",
    "\n",
    "        if name == \"open_account\":\n",
    "            result = open_account(args.get(\"name\"), args.get(\"email\"))\n",
    "            tool_messages.append({\n",
    "                \"role\": \"tool\",\n",
    "                \"tool_call_id\": tc.id,\n",
    "                \"content\": json.dumps(result)  # MUST be string\n",
    "            })\n",
    "\n",
    "        elif name == \"deposit_funds\":\n",
    "            result = deposit_funds(args[\"amount\"], new_balance)\n",
    "            new_balance = result[\"balance\"]\n",
    "            tool_messages.append({\n",
    "                \"role\": \"tool\",\n",
    "                \"tool_call_id\": tc.id,\n",
    "                \"content\": json.dumps(result)\n",
    "            })\n",
    "\n",
    "        elif name == \"withdraw_funds\":\n",
    "            result = withdraw_funds(args[\"amount\"], new_balance)\n",
    "            new_balance = result[\"balance\"]\n",
    "            tool_messages.append({\n",
    "                \"role\": \"tool\",\n",
    "                \"tool_call_id\": tc.id,\n",
    "                \"content\": json.dumps(result)\n",
    "            })\n",
    "\n",
    "        elif name == \"check_balance\":\n",
    "            result = check_balance(new_balance)\n",
    "            tool_messages.append({\n",
    "                \"role\": \"tool\",\n",
    "                \"tool_call_id\": tc.id,\n",
    "                \"content\": json.dumps(result)\n",
    "            })\n",
    "\n",
    "    return tool_messages, new_balance\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "7200d772",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "* Running on local URL:  http://127.0.0.1:7879\n",
      "* To create a public link, set `share=True` in `launch()`.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"http://127.0.0.1:7879/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def chat(message, history, balance):\n",
    "    # Build messages from history (which already includes the welcome assistant turn)\n",
    "    messages = [{\"role\": \"system\", \"content\": SYS_PROMPT}] + history + [\n",
    "        {\"role\": \"user\", \"content\": message}\n",
    "    ]\n",
    "\n",
    "    resp = openai.chat.completions.create(model=MODEL, messages=messages, tools=tools)\n",
    "    assistant_msg = resp.choices[0].message\n",
    "\n",
    "    # If the model called tools\n",
    "    if getattr(assistant_msg, \"tool_calls\", None):\n",
    "        # 1) Append the assistant's tool-call turn as a normal message dict\n",
    "        messages.append({\n",
    "            \"role\": \"assistant\",\n",
    "            \"content\": assistant_msg.content or \"\",\n",
    "            \"tool_calls\": [\n",
    "                {\n",
    "                    \"id\": tc.id,\n",
    "                    \"type\": \"function\",\n",
    "                    \"function\": {\n",
    "                        \"name\": tc.function.name,\n",
    "                        \"arguments\": tc.function.arguments,\n",
    "                    },\n",
    "                } for tc in assistant_msg.tool_calls\n",
    "            ],\n",
    "        })\n",
    "\n",
    "        # 2) Execute tools and extend with tool messages\n",
    "        tool_msgs, balance = handle_tool_call(assistant_msg, balance)\n",
    "        messages.extend(tool_msgs)  # NOT append\n",
    "\n",
    "        # 3) Ask the model to finalize using tool results\n",
    "        resp2 = openai.chat.completions.create(model=MODEL, messages=messages, tools=tools)\n",
    "        return resp2.choices[0].message.content, balance\n",
    "\n",
    "    # No tools—just return the assistant text\n",
    "    return assistant_msg.content, balance\n",
    "\n",
    "# Pre-seed the Chatbot with the welcome as the first assistant message\n",
    "chatbot = gr.Chatbot(\n",
    "    value=[{\"role\": \"assistant\", \"content\": WELCOME}],\n",
    "    type=\"messages\",\n",
    "    height=420\n",
    ")\n",
    "\n",
    "balance_state = gr.State(0.0)\n",
    "\n",
    "gr.ChatInterface(\n",
    "    fn=chat,\n",
    "    chatbot=chatbot,   # <-- inject our pre-seeded Chatbot\n",
    "    type=\"messages\",\n",
    "    additional_inputs=[balance_state],   # read current balance\n",
    "    additional_outputs=[balance_state],  # write updated balance\n",
    ").launch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98256294",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
